from flask import Flask, render_template, request, redirect, url_for, session
import time
from functools import wraps
from flask import jsonify
import socket
import threading
import json
import os
import mysql.connector
from mysql.connector import Error

# Database configuration
DB_CONFIG = {
    'host': 'localhost',
    'database': 'chat',
    'user': 'root',
    'password': 'root'
}

app = Flask(__name__)
app.secret_key = 'your_secret_key_here'

def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'username' not in session:
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    return decorated_function

def get_db_connection():
    try:
        connection = mysql.connector.connect(**DB_CONFIG)
        return connection
    except Error as e:
        print(f"Database connection failed: {e}")
        return None

def init_db():
    conn = get_db_connection()
    if conn:
        try:
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    username VARCHAR(50) NOT NULL UNIQUE,
                    password VARCHAR(255) NOT NULL,
                    locked BOOLEAN DEFAULT FALSE,
                    attempts INT DEFAULT 0,
                    lock_time TIMESTAMP NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
                )
            """)
            conn.commit()
            print("Database tables initialized")
        except Error as e:
            print(f"Failed to initialize database tables: {e}")
        finally:
            if conn.is_connected():
                cursor.close()
                conn.close()

@app.route('/', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        
        conn = get_db_connection()
        if not conn:
            return "Database connection failed"
        
        try:
            cursor = conn.cursor(dictionary=True)
            cursor.execute("SELECT * FROM users WHERE username = %s", (username,))
            user = cursor.fetchone()
            
            if not user:
                return "User not found"
            
            # Check if account is locked
            if user['locked']:
                cursor.execute("SELECT TIMESTAMPDIFF(SECOND, lock_time, NOW()) AS lock_diff FROM users WHERE username = %s", (username,))
                lock_diff = cursor.fetchone()['lock_diff']
                if lock_diff < 1800:  # 30 minutes lock
                    return "Account locked, please try again in 30 minutes"
                else:
                    cursor.execute("UPDATE users SET locked = FALSE, attempts = 0 WHERE username = %s", (username,))
                    conn.commit()
            
            # Verify password
            if user['password'] == password:
                session['username'] = username
                cursor.execute("UPDATE users SET attempts = 0 WHERE username = %s", (username,))
                conn.commit()
                return redirect(url_for('chat'))
            else:
                cursor.execute("UPDATE users SET attempts = attempts + 1 WHERE username = %s", (username,))
                cursor.execute("SELECT attempts FROM users WHERE username = %s", (username,))
                attempts = cursor.fetchone()['attempts']
                
                if attempts >= 5:
                    cursor.execute("UPDATE users SET locked = TRUE, lock_time = NOW() WHERE username = %s", (username,))
                    conn.commit()
                    return "Too many failed attempts, account locked for 30 minutes"
                
                conn.commit()
                return f"Wrong password, {5 - attempts} attempts remaining"
                
        except Error as e:
            print(f"Database operation failed: {e}")
            return "Login failed, please try again"
        finally:
            if conn.is_connected():
                cursor.close()
                conn.close()
    
    return render_template('login.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        
        conn = get_db_connection()
        if not conn:
            return "Database connection failed"
        
        try:
            cursor = conn.cursor()
            cursor.execute("INSERT INTO users (username, password) VALUES (%s, %s)", 
                          (username, password))
            conn.commit()
            return redirect(url_for('login'))
        except Error as e:
            conn.rollback()
            if "Duplicate entry" in str(e):
                return "Username already exists"
            return "Registration failed"
        finally:
            if conn.is_connected():
                cursor.close()
                conn.close()
    
    return render_template('register.html')

# Message queue and online users
message_queue = {}
online_users = set()
offline_messages = {}

# UDP message server
def udp_server():
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_socket.bind(('0.0.0.0', 9999))
    
    while True:
        data, addr = udp_socket.recvfrom(1024)
        message = json.loads(data.decode())
        
        if message['type'] == 'online':
            online_users.add(message['username'])
            if message['username'] in offline_messages:
                for msg in offline_messages[message['username']]:
                    udp_socket.sendto(json.dumps(msg).encode(), addr)
                del offline_messages[message['username']]
        elif message['type'] == 'offline':
            online_users.discard(message['username'])
        elif message['type'] == 'message':
            if message['to'] in online_users:
                pass  # Forward to online user
            else:
                if message['to'] not in offline_messages:
                    offline_messages[message['to']] = []
                offline_messages[message['to']].append(message)

# TCP file server
def tcp_server():
    tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_socket.bind(('0.0.0.0', 9998))
    tcp_socket.listen(5)
    
    while True:
        conn, addr = tcp_socket.accept()
        data = conn.recv(1024)
        file_info = json.loads(data.decode())
        filename = file_info['filename']
        filesize = file_info['filesize']
        
        with open(f'received_{filename}', 'wb') as f:
            remaining = filesize
            while remaining > 0:
                data = conn.recv(min(remaining, 4096))
                f.write(data)
                remaining -= len(data)
        
        conn.close()

# Start server threads
threading.Thread(target=udp_server, daemon=True).start()
threading.Thread(target=tcp_server, daemon=True).start()

@app.route('/chat')
@login_required
def chat():
    return render_template('chat.html')

@app.route('/get_contacts')
@login_required
def get_contacts():
    conn = get_db_connection()
    if not conn:
        return jsonify({'error': 'Database connection failed'})
    
    try:
        cursor = conn.cursor(dictionary=True)
        
        # Get all users
        cursor.execute("SELECT username FROM users")
        all_users = [row['username'] for row in cursor.fetchall()]
        
        # Online users from online_users set
        online = list(online_users)
        
        return jsonify({
            'online': online,
            'all': all_users
        })
        
    except Error as e:
        print(f"Failed to get contacts: {e}")
        return jsonify({'error': 'Failed to get contacts'})
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()

@app.route('/send_message', methods=['POST'])
@login_required
def send_message():
    data = request.json
    return jsonify({'status': 'success'})

@app.route('/send_file', methods=['POST'])
@login_required
def send_file():
    file = request.files['file']
    filename = file.filename
    file.save(os.path.join('uploads', filename))
    return jsonify({'status': 'success'})

def send_file_tcp(filename, receiver):
    try:
        receiver_addr = ('receiver_ip', 9998)
        tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        tcp_socket.connect(receiver_addr)
        
        file_info = {
            'filename': os.path.basename(filename),
            'filesize': os.path.getsize(filename),
            'sender': session['username'],
            'receiver': receiver
        }
        tcp_socket.send(json.dumps(file_info).encode())
        
        with open(filename, 'rb') as f:
            while True:
                data = f.read(4096)
                if not data:
                    break
                tcp_socket.send(data)
        
        tcp_socket.close()
        return True
    except Exception as e:
        print(f"File send failed: {e}")
        return False

@app.route('/message_history', methods=['GET'])
@login_required
def message_history():
    sender_id = request.args.get('sender_id')
    receiver_id = request.args.get('receiver_id')
    
    conn = get_db_connection()
    if not conn:
        return jsonify({'status': 'error', 'message': 'Database connection failed'})
    
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute("""
            SELECT * FROM message_contex  
            WHERE (sender_id = %s AND receiver_id = %s)
            OR (sender_id = %s AND receiver_id = %s)
            ORDER BY send_time
            """, (str(sender_id), str(receiver_id), str(receiver_id), str(sender_id)))
        
        messages = cursor.fetchall()
        return jsonify({'status': 'success', 'data': messages})
        
    except Error as e:
        print(f"Database error: {e}")
        return jsonify({'status': 'error', 'message': str(e)})
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()


if __name__ == '__main__':
    init_db()
    app.run(debug=True, port=5000)
